# -*- coding: utf-8 -*-
"""
只读本地数据绘图（不再造数据）：
- 单文件：panel_data.csv，列：tv,pv,split（split∈{train,test}，大小写不敏感）
- 双文件：train.csv / test.csv，各列：tv,pv
效果：
- 散点 + 45°参考线
- 上/右边际密度（贴轴、较薄）
- 图例左上 + R² 两行；RMSE 两行放右下
- 仅两种主色；边际与散点仅透明度不同
"""

import os
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

# ========== 文件名（按需改） ==========
SINGLE_FILE = "panel_data.csv"  # 若存在则优先使用（含 tv,pv,split）
TRAIN_FILE  = "train.csv"       # 若单文件不存在则尝试双文件
TEST_FILE   = "test.csv"

# ========== 全局字体 ==========
mpl.rcParams['font.family'] = 'Times New Roman'
mpl.rcParams['axes.unicode_minus'] = False

# ========== 可选：KDE（无 SciPy 自动退化为直方图） ==========
try:
    from scipy.stats import gaussian_kde
    _HAS_KDE = True
except Exception:
    _HAS_KDE = False

try:
    import pandas as pd
    _HAS_PANDAS = True
except Exception:
    _HAS_PANDAS = False

# ========== 配色/外观参数（只有两种主色） ==========
COLOR_TRAIN = "#4CAF50"   # 绿
COLOR_TEST  = "#F39C12"   # 橙
ALPHA_PT    = 0.95         # 散点透明度
ALPHA_MARG  = 0.25        # 边际透明度

# 边际“薄厚”和形状（KDE_BW 越小越尖，KDE_AMPL 越小越矮）
TOP_HEIGHT_RATIO  = 0.58
RIGHT_WIDTH_RATIO = 0.62
KDE_BW            = 0.22
KDE_AMPL          = 0.70

LEGEND_MARKERSCALE = 1.6   # 图例圆圈相对缩放

# ========== 计算指标 ==========
def r2_rmse(y_true, y_pred):
    y_true = np.asarray(y_true); y_pred = np.asarray(y_pred)
    resid = y_pred - y_true
    rmse = float(np.sqrt(np.mean(resid**2)))
    ss_tot = np.sum((y_true - y_true.mean())**2)
    r2 = float(1 - np.sum(resid**2) / ss_tot) if ss_tot > 0 else np.nan
    return r2, rmse

# ========== 读取本地数据 ==========
def load_local_data():
    """只从本地读取数据；返回 tv_tr, pv_tr, tv_te, pv_te。"""
    if _HAS_PANDAS:
        if os.path.exists(SINGLE_FILE):
            df = pd.read_csv(SINGLE_FILE)
            cols_lower = {c.lower(): c for c in df.columns}
            need = {"tv", "pv", "split"}
            if not need.issubset(cols_lower.keys()):
                raise ValueError(f"[ERROR] {SINGLE_FILE} 需包含列 tv,pv,split。")
            df = df.rename(columns={cols_lower["tv"]:"tv",
                                    cols_lower["pv"]:"pv",
                                    cols_lower["split"]:"split"})
            tr = df[df["split"].astype(str).str.lower().eq("train")]
            te = df[df["split"].astype(str).str.lower().eq("test")]
            if tr.empty or te.empty:
                raise ValueError("[ERROR] split 列必须含有 'train' 与 'test' 两类。")
            return tr["tv"].to_numpy(), tr["pv"].to_numpy(), te["tv"].to_numpy(), te["pv"].to_numpy()

        if os.path.exists(TRAIN_FILE) and os.path.exists(TEST_FILE):
            tr = pd.read_csv(TRAIN_FILE); te = pd.read_csv(TEST_FILE)
            for name, df in [("train", tr), ("test", te)]:
                cols_lower = {c.lower(): c for c in df.columns}
                if not {"tv","pv"}.issubset(cols_lower.keys()):
                    raise ValueError(f"[ERROR] {name}.csv 需包含列 tv,pv。")
                df.rename(columns={cols_lower["tv"]:"tv", cols_lower["pv"]:"pv"}, inplace=True)
            return tr["tv"].to_numpy(), tr["pv"].to_numpy(), te["tv"].to_numpy(), te["pv"].to_numpy()

    # 无 pandas 的兼容读取
    if os.path.exists(SINGLE_FILE):
        data = np.genfromtxt(SINGLE_FILE, delimiter=",", names=True, dtype=None, encoding="utf-8")
        if not {"tv","pv","split"}.issubset(set(data.dtype.names)):
            raise ValueError(f"[ERROR] {SINGLE_FILE} 需包含列 tv,pv,split。")
        mask_tr = np.char.lower(data["split"].astype(str)) == "train"
        mask_te = np.char.lower(data["split"].astype(str)) == "test"
        return data["tv"][mask_tr], data["pv"][mask_tr], data["tv"][mask_te], data["pv"][mask_te]

    if os.path.exists(TRAIN_FILE) and os.path.exists(TEST_FILE):
        tr = np.genfromtxt(TRAIN_FILE, delimiter=",", names=True, dtype=None, encoding="utf-8")
        te = np.genfromtxt(TEST_FILE, delimiter=",", names=True, dtype=None, encoding="utf-8")
        if not {"tv","pv"}.issubset(set(tr.dtype.names)) or not {"tv","pv"}.issubset(set(te.dtype.names)):
            raise ValueError("[ERROR] train.csv / test.csv 需包含列 tv,pv。")
        return tr["tv"], tr["pv"], te["tv"], te["pv"]

    raise FileNotFoundError(
        f"[ERROR] 未找到数据文件。请提供：\n"
        f"  方式A：{SINGLE_FILE}（列 tv,pv,split），或\n"
        f"  方式B：{TRAIN_FILE} 与 {TEST_FILE}（各含 tv,pv）。"
    )

# ========== 绘图 ==========
def plot_panel(tv_tr, pv_tr, tv_te, pv_te,
               xlab='TV of C–C/C=C', ylab='PV of C–C/C=C',
               figsize=(5.0, 5.0), dpi=200):
    from matplotlib import gridspec

    tv_tr = np.asarray(tv_tr); pv_tr = np.asarray(pv_tr)
    tv_te = np.asarray(tv_te); pv_te = np.asarray(pv_te)

    allv = np.r_[tv_tr, pv_tr, tv_te, pv_te]
    vmin, vmax = allv.min(), allv.max()
    pad = 0.04 * max(vmax - vmin, 1.0)
    lo, hi = vmin - pad, vmax + pad

    r2_tr, rmse_tr = r2_rmse(tv_tr, pv_tr)
    r2_te, rmse_te = r2_rmse(tv_te, pv_te)

    fig = plt.figure(figsize=figsize, dpi=dpi)
    gs = gridspec.GridSpec(2, 2,
                           width_ratios=[4, RIGHT_WIDTH_RATIO],
                           height_ratios=[TOP_HEIGHT_RATIO, 4],
                           wspace=0.0, hspace=0.0)

    ax_sc = fig.add_subplot(gs[1, 0])
    ax_top = fig.add_subplot(gs[0, 0], sharex=ax_sc)
    ax_rt  = fig.add_subplot(gs[1, 1], sharey=ax_sc)

    # —— 主图（四边封闭）——
    ax_sc.set_xlim(lo, hi); ax_sc.set_ylim(lo, hi)
    ax_sc.set_aspect('equal', adjustable='box')
    for s in ax_sc.spines.values(): s.set_linewidth(1.0)
    ax_sc.plot([lo, hi], [lo, hi], color='#CFCFCF', lw=1)  # y=x

    s = 12  # 散点大小
    ax_sc.scatter(tv_tr, pv_tr, s=s, c=COLOR_TRAIN, edgecolors='white',
                  linewidths=0.35, alpha=ALPHA_PT, label='Train')
    ax_sc.scatter(tv_te,  pv_te,  s=s, c=COLOR_TEST,  edgecolors='white',
                  linewidths=0.35, alpha=ALPHA_PT, label='Test')
    ax_sc.set_xlabel(xlab); ax_sc.set_ylabel(ylab)

    # 图例（左上）+ R²；RMSE 在右下
    ax_sc.legend(loc='upper left', frameon=False, fontsize=10,
                 handletextpad=0.35, borderpad=0.2, scatterpoints=1,
                 markerscale=LEGEND_MARKERSCALE)
    ax_sc.text(0.03, 0.84, f"Train R²: {r2_tr:.2f}\nTest  R²: {r2_te:.2f}",
               transform=ax_sc.transAxes, ha='left', va='top', fontsize=10)
    ax_sc.text(0.60, 0.18, f"Train RMSE: {rmse_tr:.2f}\nTest  RMSE: {rmse_te:.2f}",
               transform=ax_sc.transAxes, ha='left', va='bottom', fontsize=10)

    # —— 顶部/右侧边际（贴轴 + 更薄）——
    def kde_or_hist(ax, arr, orient='x', color=COLOR_TRAIN):
        if _HAS_KDE:
            grid = np.linspace(lo, hi, 400)
            den  = gaussian_kde(arr, bw_method=KDE_BW)(grid) * KDE_AMPL
            if orient == 'x':
                ax.fill_between(grid, 0, den, color=color, alpha=ALPHA_MARG, lw=0.6)
                ax.set_ylim(0, den.max()*1.02)
            else:
                ax.fill_betweenx(grid, 0, den, color=color, alpha=ALPHA_MARG, lw=0.6)
                ax.set_xlim(0, den.max()*1.02)
        else:
            bins = 36
            if orient == 'x':
                ax.hist(arr, bins=bins, density=True, color=color, alpha=ALPHA_MARG)
                ax.set_ylim(bottom=0)
            else:
                ax.hist(arr, bins=bins, density=True, color=color, alpha=ALPHA_MARG,
                        orientation='horizontal')
                ax.set_xlim(left=0)

    # 隐掉边框/刻度，保证与主图“无缝”
    ax_top.set_xlim(lo, hi); ax_top.axis('off')
    ax_rt.set_ylim(lo, hi);  ax_rt.axis('off')
    kde_or_hist(ax_top, tv_tr, 'x', COLOR_TRAIN)
    kde_or_hist(ax_top, tv_te,  'x', COLOR_TEST)
    kde_or_hist(ax_rt,  pv_tr, 'y', COLOR_TRAIN)
    kde_or_hist(ax_rt,  pv_te,  'y', COLOR_TEST)
    plt.savefig('test.png',dpi=300)

    plt.show()

# ========== 运行 ==========
if __name__ == "__main__":
    tv_tr, pv_tr, tv_te, pv_te = load_local_data()
    plot_panel(tv_tr, pv_tr, tv_te, pv_te,
               xlab='TV of C–C/C=C', ylab='PV of C–C/C=C')
